try: from cv2 import cv2
except: import cv2

import math
import re
import numpy as np

import matplotlib.pyplot as plt
import torch
import torch.nn as nn
from torch.autograd import Variable
from torch.utils.data import DataLoader,Dataset

class ImgItem(object):

    def __init__(self,imgname,imgpath,gt_boxes,gt_classes):
        # the img's path
        self.imgname = imgname
        self.imgpath = imgpath
        # a list of gt boxes K*4
        self.gt_boxes = list(gt_boxes)
        # a list of labels K
        self.gt_classes = list(gt_classes)
        assert len(self.gt_classes) == len(self.gt_boxes)

    def __repr__(self):
        txt = 'imgpath: {} rect: ({},{})->({},{}) class: {}'.format(self.imgpath,*self.gt_boxes[0],1)
        return txt

    def display(self):
        picture = cv2.imread(self.imgpath)
        for box in self.gt_boxes:
            picture = cv2.rectangle(picture,(box[0],box[1]),(box[2],box[3]),(0,255,0),thickness=10)
        picture = cv2.resize(picture,(255,255))
        plt.imshow(picture)
        plt.show()

class TongueDS(Dataset):

    def __init__(self,kind='train'):
        self.recordfile = '/mnt/md0/Qiu/Dataset/TouGue/imgs/result.txt'
        self.ImgItems = []
        self.W = 320
        self.H = 320
        self.kind = kind
        self.prepare_img_file()

    def prepare_img_file(self):
        with open(self.recordfile,'r') as f:
            lines = f.readlines()
            cnt = 0
            for line in lines:
                cnt += 1
                line = line.strip()
                if len(line) == 0: continue
                ret = re.findall(r'(.+).(jpg|png)：left=(.+), top=(.+), right=(.+), bottom=(.+)',line)[0]
                filename = ret[0]+'.'+ret[1]
                ret = ret[:1] + ret[2:]
                ret = list(map(int,ret))
                imgpath = '/mnt/md0/Qiu/Dataset/TouGue/imgs/'+filename
                md = cnt%5
                if self.kind=='test' and md!=0: continue
                if self.kind=='train' and md==0: continue
                self.ImgItems.append(ImgItem(ret[0],imgpath,[(ret[1],ret[2],ret[3],ret[4])],[1]))

    def __len__(self):
        return len(self.ImgItems)

    def __getitem__(self, index):
        item = self.ImgItems[index]
        img = cv2.imread(item.imgpath)
        W,H,C = img.shape
        ret_boxes = []
        for box in item.gt_boxes:
            box = list(box)
            box[0],box[2]=int(box[0]/H*self.H),int(box[2]/H*self.H)
            box[1],box[3]=int(box[1]/W*self.W),int(box[3]/W*self.W)
            ret_boxes.append(box)
        img = cv2.resize(img,(self.H,self.W))
        ret_boxes = np.array(ret_boxes).reshape(-1)
        ret_classes = np.array(item.gt_classes)
        return {'images':img,'gt_boxes':ret_boxes,'gt_classes':ret_classes}


def DLTransform(out,volatile=False):
    images = Variable(out['images'].permute(0,3,1,2),volatile=volatile).float().cuda()
    gt_boxes = out['gt_boxes']
    boxes_list = []
    for item in gt_boxes:
        boxes_list.append(item.numpy()[np.newaxis,:])
    gt_classes = out['gt_classes']
    classes_list = []
    for item in gt_classes:
        classes_list.append(item.numpy())
    empty_list = [[] for _ in range(len(boxes_list))]
    return images,boxes_list,classes_list,empty_list,0

if __name__=='__main__':

    DS = TongueDS(kind='test')
    len(DS)
    DS.ImgItems[50].display()
    item = DS.ImgItems[50]

    x = DS[50]

    picture = x['images']
    box = x['gt_boxes'][0]
    for box in x['gt_boxes']:
        picture = cv2.rectangle(picture,(box[0],box[1]),(box[2],box[3]),(0,255,0),thickness=10)
    picture = cv2.resize(picture,(255,255))
    plt.imshow(picture)
    plt.show()

    DL = DataLoader(DS,batch_size=3,shuffle=True)

    for i,out in enumerate(DL):
        break

    K = 0
    picture = out['images'][K].numpy()
    box = out['gt_boxes'][K,:].numpy()
    picture = cv2.rectangle(picture,(box[0],box[1]),(box[2],box[3]),(0,255,0),thickness=10)
    picture = cv2.resize(picture,(255,255))
    plt.imshow(picture)
    plt.show()
